# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import collections
import os
import datetime
import logging
import sys
import numpy
import pandas

from absl import app, flags

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from coin.support.proto_log.logic.util import run_from_og_log_archive
from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog, OrderEvent
import google.protobuf.json_format as json_format

FLAGS = flags.FLAGS

from coin.strategy.mm.tool.log_pta_dumper import OrderIdConverter
from coin.strategy.mm.tool.taker_ratio import autoquantile
from coin.strategy.mm.tool.order_reader import side_to_sign


class FrAnalysis(object):
  def __init__(self, market_type, exchange, trading_date, machine):
    self._trading_date = trading_date
    self._market_type = market_type
    self._exchange = exchange
    self._machine = machine
    self._subprcs = collections.defaultdict(lambda: collections.defaultdict(float))
    self._subsigns = collections.defaultdict(lambda: collections.defaultdict(float))
    self._subqtys = collections.defaultdict(lambda: collections.defaultdict(float))
    self._fillqtys = collections.defaultdict(lambda: collections.defaultdict(float))
    self._fillpqs = collections.defaultdict(lambda: collections.defaultdict(float))
    self._order_id_converter = OrderIdConverter()

  def on_log2(self, timestamp, log):
    pb = OrderGatewayLog()
    try:
      pb.ParseFromString(log)
    except Exception:
      print(log)
      return
    if pb.type == OrderGatewayLog.ORDER_EVENT:
      event = pb.event
      if event.type not in [
          OrderEvent.ORDER_SUBMITTED, OrderEvent.ORDER_ACCEPTED, OrderEvent.ORDER_FILLED
      ]:
        return

      if self._exchange == 'Bitflyer' and event.type == OrderEvent.ORDER_SUBMITTED:
        return
      elif self._exchange != 'Bitflyer' and event.type == OrderEvent.ORDER_ACCEPTED:
        return

      if self._exchange == 'Bitflyer':
        order_id = event.external_order_id
      else:
        order_id = self._order_id_converter.get_order_id(event)

      if event.type in [OrderEvent.ORDER_SUBMITTED, OrderEvent.ORDER_ACCEPTED]:
        self._subprcs[event.symbol][order_id] = event.order_price
        self._subsigns[event.symbol][order_id] = side_to_sign(event.order_side)
        self._subqtys[event.symbol][order_id] = event.order_qty
      elif event.type == OrderEvent.ORDER_FILLED:
        if self._exchange == 'Bitflyer':
          self._fillqtys[event.symbol][order_id] = event.fill_qty
          self._fillpqs[event.symbol][order_id] = event.fill_qty * event.fill_price
        else:
          self._fillqtys[event.symbol][order_id] += event.fill_qty
          self._fillpqs[event.symbol][order_id] += event.fill_qty * event.fill_price

  def plot_dicts(self, symbol, subqtys, fillqtys, fillpqs, subprcs, subsigns):
    tdstr = self._trading_date.strftime("%Y%m%d")
    nametuple = (self._exchange, self._machine, symbol, tdstr)

    subqs = []
    fillqs = []
    fillrs = []

    for iid in subqtys.keys():
      subqs.append(subqtys[iid])
      fillqs.append(fillqtys[iid])
      fillrs.append(fillqtys[iid] / subqtys[iid])

    freturns = []
    returns = []
    next_mtmps = []
    for iid in reversed(list(subqtys.keys())):
      if iid in fillpqs:
        fillprc = fillpqs[iid] / fillqtys[iid]
      else:
        fillprc = None
      subprc = subprcs[iid]
      subsign = subsigns[iid]
      next_mtmps.append(subprc)
      if subsign > 0:
        ret30 = 0 if subprc is None else numpy.mean(next_mtmps[-30:]) - subprc
        fillret30 = 0 if fillprc is None else numpy.mean(next_mtmps[-30:]) - fillprc
      else:
        ret30 = 0 if subprc is None else subprc - numpy.mean(next_mtmps[-30:])
        fillret30 = 0 if fillprc is None else fillprc - numpy.mean(next_mtmps[-30:])
      returns.append(ret30 / subprc)
      freturns.append(fillret30 / subprc)

    returns = list(reversed(returns))
    freturns = list(reversed(freturns))

    filld = pandas.DataFrame({
        'subqs': subqs, 'fillqs': fillqs, 'returns': returns, 'freturns': freturns
    }).groupby('subqs').mean().to_dict()

    fqd = filld['fillqs']
    frd = filld['returns']
    ffrd = filld['freturns']

    fillqms = []
    returnms = []
    freturnms = []
    for i, subq in enumerate(subqs):
      fillqms.append(fqd[subq])
      returnms.append(frd[subq])
      freturnms.append(ffrd[subq])

    sq, fq = autoquantile(subqs, fillqms, 5)
    sq = numpy.array(sq)

    xx = numpy.vstack([subqs, numpy.ones(len(subqs))]).T
    yy = fillrs

    aa, _, _, _ = numpy.linalg.lstsq(xx, yy)
    plt.subplot(311)
    plt.plot(sq, fq, 'g.-', linewidth=0.5)
    plt.plot(sq, aa[0] * sq + aa[1], 'r--', linewidth=0.5)
    namestr = "_".join(nametuple)
    plt.title(namestr)
    plt.subplot(312)
    plt.plot(subqs, fillrs, 'g.', markersize=0.5)
    plt.subplot(313)
    modeq = collections.Counter(subqs).most_common()[0][0]
    fqdist = sorted([fillq for subq, fillq in zip(subqs, fillrs) if subq == modeq])
    plt.plot(fqdist, 'g--', linewidth=0.5)
    plt.title("fill ratio dist when ls==%s" % modeq)
    plt.savefig("%s/fr_%s.png" % (flags.FLAGS.out_dir, namestr))
    plt.close()

    sq, rq = autoquantile(subqs, returnms, 5)
    sq, frq = autoquantile(subqs, freturnms, 5)

    # plt.subplot(211)
    # plt.plot(sq, rq, 'g.-', sq, frq, 'r.-', linewidth=0.5)
    # plt.axhline(y=2e-4, color='k', linestyle='--', lw=0.5)
    # plt.title("%s %s subq vs ret (fillret blue)" %  (tdstr, symbol))
    # plt.subplot(212)
    plt.title("%s %s fqdist (fr: %.2f)" % (tdstr, symbol, sum(fillqs) / sum(subqs)))
    plt.plot(sorted(fillqs), 'g--', linewidth=0.5)
    plt.tight_layout()
    plt.savefig("%s/ret_%s.png" % (flags.FLAGS.out_dir, namestr))
    plt.close()

    return {
        'fill_ratio': pandas.Series(fillqs).sum() / pandas.Series(subqs).sum(),
        'fill_msg_ratio': (pandas.Series(fillqs) > 0).sum() / (pandas.Series(subqs) > 0).sum(),
        'avg_return': pandas.Series(freturnms).mean()
    }

  def export_log(self):
    summary = {}
    for symbol in self._subqtys.keys():
      if self._exchange in ('Okex', 'Huobi') and self._market_type == 'Futures':
        if self._exchange == 'Okex':
          from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
          try:
            product = OkexFuturesProduct.FromStr(symbol, self._trading_date)
            repr_symbol = product.subscription_symbol
          except Exception:
            continue
        else:
          from coin.exchange.huobi_futures.kr_rest.futures_product import HuobiFuturesProduct
          try:
            product = HuobiFuturesProduct.FromStr(symbol, self._trading_date)
            repr_symbol = product.subscription_symbol
          except Exception:
            continue
      else:
        repr_symbol = symbol
      summary[repr_symbol] = self.plot_dicts(repr_symbol,
                                             self._subqtys[symbol],
                                             self._fillqtys[symbol],
                                             self._fillpqs[symbol],
                                             self._subprcs[symbol],
                                             self._subsigns[symbol])
    return summary


def launch(exchange, trading_date, hours, root_dir, machine):
  start_time = trading_date  # UTC 0 hour
  # 24 will omit some logs..

  strategy = FrAnalysis(FLAGS.market_type, exchange, trading_date, machine)
  logging.info("start log analysis")
  run_from_og_log_archive(strategy.on_log2,
                          FLAGS.market_type,
                          exchange,
                          start_time,
                          start_time + datetime.timedelta(hours=hours),
                          root_dir,
                          machine,
                          owner=None)

  return strategy.export_log()


def get_trading_dates(trading_date_str):
  if trading_date_str.find("-") >= 0:
    tdfrom, tdto = trading_date_str.split("-")
    return pandas.date_range(tdfrom, tdto).to_pydatetime()
  else:
    return [datetime.datetime.strptime(trading_date_str, '%Y%m%d')]


def main(_):
  # import warnings; warnings.filterwarnings("error")
  exchange = FLAGS.exchange
  trading_date = FLAGS.trading_date
  root_dir = os.path.expanduser(FLAGS.root_dir)
  machine = FLAGS.machine
  assert FLAGS.market_type in ('ExWallet', 'Spot', 'Margin', 'Futures')
  assert trading_date, '--trading_date must be specified.'
  assert machine, '--machine must be specified.'
  assert exchange, '--exchange must be specified.'
  print('Running for %s %s %s ...' % (trading_date, machine, exchange))
  sys.stdout.flush()
  trading_dates = get_trading_dates(trading_date)
  if len(trading_dates) == 1:
    launch(exchange, trading_dates[0], 25, root_dir, machine)
  else:
    summary = collections.defaultdict(
        lambda: collections.defaultdict(lambda: collections.OrderedDict()))
    for trading_date in trading_dates:
      day_summary = launch(exchange, trading_date, 25, root_dir, machine)
      for symbol in day_summary:
        for field in day_summary[symbol]:
          summary[symbol][field][trading_date] = day_summary[symbol][field]
    for symbol in summary:
      df = pandas.DataFrame(summary[symbol])
      for plttype in ["fr_qty", "fr_msg", "avg_ret"]:
        if plttype == "fr_qty":
          plt.plot(df['fill_ratio'], linewidth=0.5)
        elif plttype == "fr_msg":
          plt.plot(df['fill_msg_ratio'], linewidth=0.5)
        else:
          plt.plot(df['avg_return'], 'g.-', linewidth=0.5)
        plt.title(f'{symbol} {plttype}')
        plt.xticks(rotation=20)
        plt.savefig("%s/summ_%s_%s_%s.png" % (flags.FLAGS.out_dir, machine, symbol, plttype))
        plt.close()


if __name__ == '__main__':
  flags.DEFINE_string('exchange', 'Okex', 'Exchange name.')

  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('root_dir', '~/data/og_proto_log', 'root_dir.')

  flags.DEFINE_string('machine', 'strategy-04.aliyun-cn-hongkong', 'Instance machine name.')

  flags.DEFINE_string('market_type', 'Futures', 'Spot, Futures, Margin')

  flags.DEFINE_string('out_dir', 'oe_stat_fr', 'out_directory')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
